import cv2
import numpy as np
import os
from os.path import exists
from imutils import paths
import pickle
from tqdm import tqdm
from tensorflow.keras.applications.vgg16 import VGG16, preprocess_input
from tensorflow.keras.preprocessing import image
import logging

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# 定义函数获取文件大小
def get_size(file):
    return os.path.getsize(file) / (1024 * 1024)

# 定义函数createXY
def createXY(train_folder, dest_folder, method='vgg', batch_size=64):
    x_file_path = os.path.join(dest_folder, "X.pkl")
    y_file_path = os.path.join(dest_folder, "y.pkl")

    # 确保目标目录存在
    if not os.path.exists(dest_folder):
        try:
            os.makedirs(dest_folder)
            logging.info(f"创建目录：{dest_folder}")
        except OSError as e:
            logging.error(f"创建目录 {dest_folder} 时发生错误：{e}")
            return None, None

    # 如果 X 和 y 已经存在，则直接读取，不再重新构建
    if os.path.exists(x_file_path) and os.path.exists(y_file_path):
        logging.info("X和y已经存在，直接读取")
        logging.info(f"X文件大小:{get_size(x_file_path):.2f}MB")
        logging.info(f"y文件大小:{get_size(y_file_path):.2f}MB")
        
        with open(x_file_path, 'rb') as f:
            X = pickle.load(f)
        with open(y_file_path, 'rb') as f:
            y = pickle.load(f)
        
        return X, y

    logging.info("读取所有图像，生成X和y")
    image_paths = list(paths.list_images(train_folder))

    X = []
    y = []
    # 根据传入的方法选择不同的模型
    if method == 'vgg':
        model = VGG16(weights='imagenet', include_top=False, pooling="max")
        logging.info("完成构建 VGG16 模型")
    elif method == 'flat':
        model = None

    # 计算需要的批次数
    num_batches = len(image_paths) // batch_size + (1 if len(image_paths) % batch_size else 0)

    # 使用进度条对批次进行循环处理
    for idx in tqdm(range(num_batches), desc="读取图像"):
        batch_images = []
        batch_labels = []

        start = idx * batch_size
        end = min((idx + 1) * batch_size, len(image_paths))

        # 对于每个批次中的图像
        for i in range(start, end):
            image_path = image_paths[i]
            try:
                if method == 'vgg':
                    img = image.load_img(image_path, target_size=(224, 224))
                    img = image.img_to_array(img)
                elif method == 'flat':
                    img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
                    img = cv2.resize(img, (32, 32))
                
                batch_images.append(img)
                
                label = os.path.basename(image_path).split(".")[0]
                label = 1 if 'dog' in label else 0
                batch_labels.append(label)
            except Exception as e:
                logging.error(f"处理图像 {image_path} 时发生错误：{e}")

        batch_images = np.array(batch_images)
        if method == 'vgg':
            batch_images = preprocess_input(batch_images)
            batch_pixels = model.predict(batch_images, verbose=0)
        else:
            batch_pixels = batch_images.reshape(batch_images.shape[0], -1)
        
        X.extend(batch_pixels)
        y.extend(batch_labels)

    logging.info(f"X.shape: {np.shape(X)}")
    logging.info(f"y.shape: {np.shape(y)}")
        
    # 将X和y分别序列化到文件
    with open(x_file_path, 'wb') as f:
        pickle.dump(X, f)
    with open(y_file_path, 'wb') as f:
        pickle.dump(y, f)

    return X, y